import torch
import sys
import os
sys.path.append(os.path.dirname(__file__))

from lib.utils.learning import *
from lib.utils.tools import get_config
from lib.utils.utils_data import flip_data
from lib.data.dataset_wild import read_input

import os
import collections

def clean_state_dict(state_dict):
    new_state_dict = collections.OrderedDict()
    for k, v in state_dict.items():
        if k[:7] == 'module.':
            k = k[7:]  # remove `module.`
        new_state_dict[k] = v
    return new_state_dict 
# 原文链接：https://blog.csdn.net/hxxjxw/article/details/119491163

class motionbert_inferencer:

    def __init__(self, config_path, ckpt_path):
        # breakpoint()
        self.args = get_config(config_path)
        self.model = load_backbone(self.args)

        self.ckpt_path = ckpt_path
        self.checkpoint = torch.load(self.ckpt_path, map_location=lambda storage, loc: storage)   
        """
        (Pdb) type(self.checkpoint)
        <class 'dict'>
        (Pdb) self.checkpoint.keys()
        dict_keys(['epoch', 'lr', 'optimizer', 'model_pos', 'min_loss'])
        (Pdb) type(self.checkpoint['model_pos'])
        <class 'collections.OrderedDict'>
        (Pdb) self.checkpoint['model_pos'].keys()
        odict_keys(['module.temp_embed', 'module.pos_embed', 'module.joints_embed.weight', ...

        """
        self.checkpoint['model_pos'] = clean_state_dict(self.checkpoint['model_pos'])
        self.model.load_state_dict(self.checkpoint['model_pos'], strict=True)
        self.model.reconstruct_temp_embed(2000)
        
        if torch.cuda.is_available():
            self.model = nn.DataParallel(self.model)
            self.model = self.model.cuda()
        self.model.eval()

    def __call__(self, json_path, block_lower=False):
        pose_data = torch.from_numpy(read_input(json_path, None, [1, 1], None))
        pose_data = torch.unsqueeze(pose_data, dim=0)

        results_all = []
        with torch.no_grad():
            if torch.cuda.is_available():
                pose_data = pose_data.cuda()
            if self.args.no_conf:
                pose_data = pose_data[:, :, :, :2]
            if self.args.flip:    
                pose_data_flip = flip_data(pose_data)
                predicted_3d_pos_1 = self.model(pose_data, block_lower=block_lower)
                predicted_3d_pos_flip = self.model(pose_data_flip, block_lower=block_lower)
                predicted_3d_pos_2 = flip_data(predicted_3d_pos_flip) # Flip back
                predicted_3d_pos = (predicted_3d_pos_1 + predicted_3d_pos_2) / 2.0
            else:
                predicted_3d_pos = self.model(pose_data)
            if self.args.rootrel:
                predicted_3d_pos[:,:,0,:]=0                    # [N,T,17,3]
            else:
                predicted_3d_pos[:,0,0,2]=0
                pass
            if self.args.gt_2d:
                predicted_3d_pos[...,:2] = pose_data[...,:2]
            results_all.append(predicted_3d_pos.cpu().numpy())
        results_all = np.hstack(results_all)
        results_all = np.concatenate(results_all)

        return results_all      

if __name__ == "__main__":
    config_path = os.path.join(os.path.dirname(__file__), "configs/pose3d/MB_train_h36m.yaml")
    ckpt_path = os.path.join(os.path.dirname(__file__), 
                             "checkpoint/pose3d/MB_train_h36m_dual_loss/finish_at_2024_4_14/best_epoch.bin")
    motionbert_inferencer(config_path, ckpt_path)
